import os
from typing import Tuple

import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import savgol_filter
import MDAnalysis as mda

import GraphTools.frame_series_method as fm
import GraphTools.graph_method as gm


def analyze_mechanical_properties(sub_mission_path: str) -> None:
    """
    Analyze mechanical properties for a given sub-mission directory.

    This function performs the following steps:
      1. Creates an output folder 'analysis_基础力学性质' under sub_mission_path.
      2. Finds and loads the first XYZ file to build an MDAnalysis Universe.
      3. Determines the fracture frame and corresponding break strain.
      4. Loads the stress-strain data and computes material toughness up to break strain.
      5. Applies Savitzky-Golay smoothing to the stress data and finds the maximum stress point.
      6. Saves the stress-strain plot and writes key results into a text file.

    :param sub_mission_path: Path to the sub-mission directory containing simulation outputs.
    """
    output_dir = os.path.join(sub_mission_path, 'analysis')
    os.makedirs(output_dir, exist_ok=True)

    # Step 1: Locate and load the XYZ file
    xyz_file = _find_xyz_file(sub_mission_path)
    universe = mda.Universe(xyz_file, format='LAMMPSDUMP')

    # Step 2: Identify fracture frame and break strain
    broken_frame, n_frames, break_location = gm.find_broken_frame(universe)
    times, strains = fm.universe_strain_time_series(universe)
    break_strain = strains[broken_frame]
    break_times = times[broken_frame]

    # Step 3: Load stress-strain data
    strain, stress = _load_strain_stress(sub_mission_path)

    # Step 4: Compute toughness (area under curve up to break strain)
    toughness = _compute_toughness(strain, stress, break_strain)

    # Step 5: Smooth stress curve and find maximum stress point
    break_idx = _find_break_index(strain, break_strain)
    strain_cut = strain[: break_idx + 1]
    stress_cut = stress[: break_idx + 1]
    smoothed = savgol_filter(stress_cut, window_length=21, polyorder=3)
    max_idx = int(np.argmax(smoothed))
    max_strain = float(strain_cut[max_idx])
    max_stress = float(smoothed[max_idx])

    # Step 6: Save the plot to the output directory
    plot_path = os.path.join(output_dir, 'stress_strain_curve.png')
    _plot_curve(strain_cut, stress_cut, smoothed,
                max_strain, max_stress, plot_path)

    # Step 7: Record results in a text file
    result_file = os.path.join(output_dir, 'analysis_基础力学性质.txt')
    with open(result_file, 'w', encoding='utf-8') as f:
        f.write("名称\t数据\n")
        f.write(f"材料韧性\t{toughness:.3f}\n")
        f.write(f"断裂应变\t{break_strain:.3f}\n")
        f.write(f"最大应变\t{max_strain:.3f}\n")
        f.write(f"最大应力\t{max_stress:.3f}\n")
        f.write(f"断裂帧\t{broken_frame}\n")
        f.write(f"断裂时间\t{break_times}\n")
        f.write(f"总帧数\t{n_frames}\n")
        f.write(f"断裂位置\t{break_location}\n")


# ----- Helper functions -----

def _find_xyz_file(path: str) -> str:
    """Find the first LAMMPS dump '.xyz' file in the given directory."""
    for fname in os.listdir(path):
        if fname.startswith('out') and fname.endswith('.xyz'):
            return os.path.join(path, fname)
    raise FileNotFoundError(f"未找到以 'out*.xyz' 命名的文件于: {path}")


def _load_strain_stress(path: str) -> Tuple[np.ndarray, np.ndarray]:
    """Load strain and stress data from the analysis text file."""
    txt_path = os.path.join(path, 'analysis', 'v_strain_xv_strees_x_.txt')
    if not os.path.isfile(txt_path):
        raise FileNotFoundError(f"应力-应变数据文件不存在: {txt_path}")
    data = np.loadtxt(txt_path, skiprows=1)
    return data[:, 0], data[:, 1]


def _compute_toughness(strain: np.ndarray, stress: np.ndarray, break_strain: float) -> float:
    """Compute the area under the stress-strain curve up to the break strain."""
    idx = _find_break_index(strain, break_strain)
    return float(np.trapz(stress[: idx + 1], strain[: idx + 1]))


def _find_break_index(strain: np.ndarray, break_strain: float) -> int:
    """Return the last index where strain <= break_strain."""
    indices = np.where(strain <= break_strain)[0]
    if not len(indices):
        raise ValueError(f"断裂应变 {break_strain:.3f} 不在应变数据范围内.")
    return int(indices[-1])


def _plot_curve(
    strain: np.ndarray,
    stress: np.ndarray,
    smoothed: np.ndarray,
    max_strain: float,
    max_stress: float,
    save_path: str,
) -> None:
    """Plot the raw and smoothed stress-strain curves, mark max stress point, and save image."""
    plt.figure(figsize=(10, 6))
    plt.plot(strain, stress, 'b-', alpha=0.5, label='Raw Data')
    plt.plot(strain, smoothed, 'r-', label='Smoothed Data')
    plt.plot(max_strain, max_stress, 'ko', label='Max Stress Point')
    plt.annotate(
        f'({max_strain:.3f}, {max_stress:.3f})',
        xy=(max_strain, max_stress),
        xytext=(10, 10),
        textcoords='offset points'
    )
    plt.xlabel('Strain')
    plt.ylabel('Stress')
    plt.title('Stress-Strain Curve (Raw vs Smoothed)')
    plt.legend()
    plt.grid(True)
    plt.savefig(save_path)
    plt.close()


def read_breaking_points_from_files(multi_massion_space):
    """
    从 multi_massion_space 下所有实验组的 analysis_基础力学性质.txt 文件中读取断裂点信息。
    
    参数:
        multi_massion_space (str): 包含多个实验组目录的路径
        
    返回:
        tuple: (broken_frames, avg_broken_frame)
            - broken_frames: 包含所有实验组断裂点比例的列表
            - avg_broken_frame: 所有断裂点比例的平均值
    """
    broken_frames = []  # 存放每个任务中得到的断裂点比例
    
    # 遍历所有子文件夹
    for root, dirs, files in os.walk(multi_massion_space):
        # 检查当前目录是否为实验组目录(N0-N9)
        if os.path.basename(root).startswith('N') and os.path.basename(root)[1:].isdigit():
            # 构建 analysis_基础力学性质.txt 的路径
            analysis_file = os.path.join(root, 'analysis', 'analysis_基础力学性质.txt')
            
            if os.path.exists(analysis_file):
                try:
                    # 读取文件内容
                    with open(analysis_file, 'r') as f:
                        lines = f.readlines()
                        
                    # 解析文件内容
                    broken_frame = None
                    n_frames = None
                    
                    for line in lines[1:]:  # 跳过标题行
                        fields = line.strip().split('\t')
                        if len(fields) >= 2:
                            key = fields[0].strip()
                            value = fields[1].strip()
                            
                            if key == '断裂帧':
                                try:
                                    broken_frame = float(value)
                                except ValueError:
                                    print(f"无法将断裂帧值转换为数字: {value}")
                            elif key == '总帧数':
                                try:
                                    n_frames = float(value)
                                except ValueError:
                                    print(f"无法将总帧数值转换为数字: {value}")
                    
                    # 如果找到了两个值，计算比例
                    if broken_frame is not None and n_frames is not None and n_frames > 0:
                        broken_frames.append(broken_frame / n_frames)
                    else:
                        print(f"在文件 {analysis_file} 中未找到完整的断裂帧和总帧数信息")
                            
                except Exception as e:
                    print(f"处理文件 {analysis_file} 时出错: {e}")
    
    # 计算所有断裂点比例的平均值
    if broken_frames:
        avg_broken_frame = np.mean(broken_frames)
    else:
        avg_broken_frame = None
        
    return broken_frames, avg_broken_frame

def process_mechanical_properties(multi_massion_space, mode='range'):
    """
    处理力学性质数据并生成图表，包括应力-应变曲线和断键-应变曲线。
    
    参数:
        multi_massion_space (str): 包含多个实验组目录的路径
        mode (str): 绘图模式，'std' 或 'range'，默认为 'range'
        
    返回:
        tuple: (results_dir, success)
            - results_dir: 结果保存的目录路径
            - success: 是否成功处理数据
    """
    import os
    import re
    import numpy as np
    import pandas as pd
    import matplotlib.pyplot as plt
    from GraphTools import analysis, basic_method as bm
    
    # 创建结果保存目录
    results_dir = os.path.join(multi_massion_space, 'results')
    os.makedirs(results_dir, exist_ok=True)
    
    # 1. 处理应力-应变数据
    def process_stress_strain():
        x_data = None
        data_list = []
        x_label = "Strain"
        y_label = "Stress"
        
        pattern = re.compile(r'^v_strain_xv_strees_x_.*\.txt$')
        target_dirs = bm.collect_target_dirs(multi_massion_space)
        
        if not target_dirs:
            print("没有找到满足文件字段一致性要求的目标文件夹。")
            return None, None, x_label, y_label
            
        for folder in sorted(target_dirs):
            analysis_path = os.path.join(folder, 'analysis')
            if os.path.isdir(analysis_path):
                for file in sorted(os.listdir(analysis_path)):
                    if pattern.match(file):
                        file_path = os.path.join(analysis_path, file)
                        try:
                            df = pd.read_csv(file_path, sep='\t', engine='python')
                        except Exception as e:
                            print(f"读取文件 {file_path} 时出错: {e}")
                            continue
                            
                        if df.shape[1] < 2:
                            continue
                            
                        cur_x = df.iloc[:, 0].values
                        cur_y = df.iloc[:, 1].values
                        
                        if x_data is None:
                            x_data = cur_x
                        elif not np.allclose(x_data, cur_x, rtol=1e-5, atol=1e-8):
                            print(f"警告: 文件 {file_path} 的 X 数据与之前文件不一致！")
                            continue
                            
                        data_list.append(cur_y)
                        break
                        
        if x_data is None or len(data_list) == 0:
            return None, None, x_label, y_label
            
        data_array = np.column_stack(data_list)
        data_df = pd.DataFrame(data_array, columns=[f"file_{i+1}" for i in range(data_array.shape[1])])
        
        return x_data, data_df, x_label, y_label
    
    # 2. 处理断键-应变数据
    def process_broken_bonds():
        x_y_list = []
        x_label = "Strain"
        y_label = "Broken_Bonds"
        
        pattern = re.compile(r'^out_Strain—Broken_Bonds.*\.txt$')
        target_dirs = bm.collect_target_dirs(multi_massion_space)
        
        if not target_dirs:
            return None, None, None, None
            
        for folder in sorted(target_dirs):
            analysis_path = os.path.join(folder, 'analysis')
            if os.path.isdir(analysis_path):
                for file in sorted(os.listdir(analysis_path)):
                    if pattern.match(file):
                        file_path = os.path.join(analysis_path, file)
                        try:
                            df = pd.read_csv(file_path, sep='\t', engine='python')
                        except Exception as e:
                            print(f"读取文件 {file_path} 时出错: {e}")
                            continue
                            
                        if df.shape[1] < 2:
                            continue
                            
                        cur_x = df.iloc[:, 0].values
                        cur_y = df.iloc[:, 1].values
                        
                        if len(x_y_list) == 0:
                            x_label = df.columns[0]
                            y_label = df.columns[1]
                            
                        x_y_list.append((cur_x, cur_y))
                        break
                        
        if not x_y_list:
            return None, None, None, None
            
        # 插值处理
        min_x = min(np.min(x) for x, _ in x_y_list)
        max_x = max(np.max(x) for x, _ in x_y_list)
        X_common = np.linspace(min_x, max_x, 200)
        
        Y_common_list = []
        for cur_x, cur_y in x_y_list:
            y_interp = np.interp(X_common, cur_x, cur_y)
            Y_common_list.append(y_interp)
            
        data_array = np.column_stack(Y_common_list)
        data_df = pd.DataFrame(data_array, columns=[f"file_{i+1}" for i in range(data_array.shape[1])])
        
        return X_common, data_df, x_label, y_label
    
    # 3. 读取断裂点信息
    broken_frames, avg_broken_frame = read_breaking_points_from_files(multi_massion_space)
    
    # 4. 处理并绘制应力-应变曲线
    x_data_ss, data_df_ss, x_label_ss, y_label_ss = process_stress_strain()
    if x_data_ss is not None and data_df_ss is not None:
        plt_ss = analysis.plot_with_shadow(
            x_data=x_data_ss,
            data_df=data_df_ss,
            mode=mode,
            note='Strain-Stress',
            x_label=x_label_ss,
            y_label=y_label_ss
        )
        
        # 添加断裂点标记
        if broken_frames and avg_broken_frame is not None:
            ax = plt_ss.gca()
            x_min, x_max = x_data_ss.min(), x_data_ss.max()
            
            region_x_min = x_min + min(broken_frames) * (x_max - x_min)
            region_x_max = x_min + max(broken_frames) * (x_max - x_min)
            broken_x_coords = [x_min + b * (x_max - x_min) for b in broken_frames]
            avg_x = x_min + avg_broken_frame * (x_max - x_min)
            
            ax.axvspan(region_x_min, region_x_max, color='yellow', alpha=0.2, label='Broken Frames Range')
            for x_coord in broken_x_coords:
                ax.axvline(x_coord, color='orange', linestyle='--', alpha=0.5)
            ax.axvline(avg_x, color='red', linewidth=2.5, label='Average Broken Frame')
            ax.legend()
            
        plt_ss.savefig(os.path.join(results_dir, f'Strain-Stress_shadow_mode_{mode}.png'))
        plt.close(plt_ss)
    
    # 5. 处理并绘制断键-应变曲线
    x_data_bb, data_df_bb, x_label_bb, y_label_bb = process_broken_bonds()
    if x_data_bb is not None and data_df_bb is not None:
        plt_bb = analysis.plot_with_shadow(
            x_data=x_data_bb,
            data_df=data_df_bb,
            mode=mode,
            note='Strain-Broken_Bonds',
            x_label=x_label_bb,
            y_label=y_label_bb
        )
        
        # 添加断裂点标记
        if broken_frames and avg_broken_frame is not None:
            ax = plt_bb.gca()
            x_min, x_max = x_data_bb.min(), x_data_bb.max()
            
            region_x_min = x_min + min(broken_frames) * (x_max - x_min)
            region_x_max = x_min + max(broken_frames) * (x_max - x_min)
            broken_x_coords = [x_min + b * (x_max - x_min) for b in broken_frames]
            avg_x = x_min + avg_broken_frame * (x_max - x_min)
            
            ax.axvspan(region_x_min, region_x_max, color='yellow', alpha=0.2, label='Broken Frames Range')
            for x_coord in broken_x_coords:
                ax.axvline(x_coord, color='orange', linestyle='--', alpha=0.5)
            ax.axvline(avg_x, color='red', linewidth=2.5, label='Average Broken Frame')
            ax.legend()
            
        plt_bb.savefig(os.path.join(results_dir, f'Strain-Broken_Bonds_shadow_mode_{mode}.png'))
        plt.close(plt_bb)
    # 保存应力-应变曲线的原始数据
    if x_data_ss is not None and data_df_ss is not None:
        data_ss = pd.DataFrame({
            'Strain': x_data_ss,
            **data_df_ss
        })
        data_ss.to_csv(os.path.join(results_dir, f'Strain-Stress_shadow_mode_{mode}.txt'), 
                       sep='\t', index=False)
    
    # 保存断键-应变曲线的原始数据
    if x_data_bb is not None and data_df_bb is not None:
        data_bb = pd.DataFrame({
            'Strain': x_data_bb,
            **data_df_bb
        })
        data_bb.to_csv(os.path.join(results_dir, f'Strain-Broken_Bonds_shadow_mode_{mode}.txt'),
                       sep='\t', index=False)

    success = (x_data_ss is not None and data_df_ss is not None) or \
              (x_data_bb is not None and data_df_bb is not None)
              
    return results_dir, success
